package main

import (
	"database/sql"
	"flag"
	"fmt"
	"runtime"
	"sync"
	"sync/atomic"
	"syscall"
	"time"

	_ "github.com/mattn/go-sqlite3"
	"github.com/vrecan/death"
)

var (
	createTable = `
		CREATE TABLE IF NOT EXISTS "keys" (
		   "id" INTEGER PRIMARY KEY AUTOINCREMENT,
		   "address" VARCHAR(42) NULL,
		   "privateKey" VARCHAR(66) NULL
		);`
	dot    = flag.Bool("dot", false, "polkadot")
	level  = flag.Int("level", 0, "beautiful level")
	export = flag.Bool("export", false, "export addresses to file")
	getkey = flag.String("getkey", "", "get address private key")
)

func init() {
	flag.Parse()
}

func main() {
	// 初始数据库
	db, err := sql.Open("sqlite3", "keys.db")
	if err != nil {
		panic(err)
	}
	defer db.Close()

	if _, err = db.Exec(createTable); err != nil {
		panic(err)
	}

	// 导出地址
	if *export {
		err = exportAddresses(db, ".", 1000000)
		if err != nil {
			panic(err)
		}
		return
	}

	stmt, err := db.Prepare("INSERT INTO keys(address, privateKey) values(?, ?);")
	if err != nil {
		panic(err)
	}
	defer stmt.Close()

	// 获取地址私钥
	if len(*getkey) > 0 {
		privateKey, err := getPrivateKey(db, *getkey)
		if err != nil {
			panic(err)
		}
		fmt.Println("address:", *getkey, "", "privateKey:", privateKey)
		return
	}

	// 扫描靓号地址
	if *level > 0 {
		addresses, err := scanBeautifulAddress(db, *level)
		if err != nil {
			panic(err)
		}
		if len(addresses) > 0 {
			fmt.Println("Beautiful Address:")
			for _, address := range addresses {
				fmt.Println(address)
			}
		}
		return
	}

	// 批量生成账号
	var count int32
	var stopped int32
	var mutex sync.Mutex
	w := sync.WaitGroup{}
	fmt.Println("Start batch generate account")
	for i := 0; i < runtime.NumCPU(); i++ {
		w.Add(1)
		go func() {
			for {
				if atomic.LoadInt32(&stopped) == 1 {
					break
				}

				var address, privateKey string
				if *dot {
					address, privateKey, err = generatePolkadotKeyPair()
				} else {
					address, privateKey, err = generateEthereumKeyPair()
				}
				if err != nil {
					panic(err)
				}

				mutex.Lock()
				if _, err = stmt.Exec(address, privateKey); err != nil {
					panic(err)
				}
				mutex.Unlock()

				if val := increment(&count); val%1000 == 0 {
					fmt.Println("time:", formatTimestamp(time.Now()), "", "keys count:", val)
				}
			}

			w.Done()
		}()
	}

	// 捕捉退出信号
	go func() {
		d := death.NewDeath(syscall.SIGTERM, syscall.SIGINT, syscall.SIGKILL, syscall.SIGQUIT)
		d.WaitForDeathWithFunc(func() {
			atomic.StoreInt32(&stopped, 1)
		})
	}()

	w.Wait()
}
